/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: TbeGatherV2的头文件
 * Author: zhangjian
 * Create: 2023-5-17
 */

#ifndef TBE_GATHERV2_H
#define TBE_GATHERV2_H

#include <string>
#include <unordered_map>
#include <nlohmann/json.hpp>
#include <hccl/hccl_types.h>
#include <mutex>
#include <map>
#include <list>
#include <hccl/base.h>
#include "tbe_vector_reduce.h"
#ifndef HCCD
#include "runtime/rt.h"
#include "graph/operator.h"
#include "graph/compute_graph.h"
#include "graph/node.h"
#endif
#include "op_tiling.h"

namespace TbeReduce {

const std::unordered_map<std::string, std::string> GATHERV2_SOC_STRING_MAP = {
    {"Ascend910B1", "_B1B2"},
    {"Ascend910B2", "_B1B2"},
    {"Ascend910B2C", "_B1B2"},
    {"Ascend910B3", "_B3B4"},
    {"Ascend910B4", "_B3B4"},
    {"Ascend910B4-1", "_B3B4-1"},
    {"Ascend910_9391", "_B1B2"},
    {"Ascend910_9381", "_B1B2"},
    {"Ascend910_9392", "_B1B2"},    // Ascend910_9392、Ascend910_9382为预留类型，当前版本暂不支持，待跟随后续版本节奏交付
    {"Ascend910_9382", "_B1B2"},
    {"Ascend910_9372", "_B3B4"},
    {"Ascend910_9362", "_B3B4-1"}
};

class TbeGatherV2 : public TbeVectorReduce {
public:
    TbeGatherV2();
    ~TbeGatherV2() override;
    HcclResult Init();
    HcclResult Run(const void *embVal, const void *refMat, u64 inputCount,
        u64 outputCount, u32 valueDim, const HcclDataType dataType, void *stream, const void *dst);
private:
    void InitOpInfoMap910A();
    void InitOpInfoMap310P3();
    void InitOpInfoMap910B();
    bool SupportRt2();
    HcclResult GetChipName(std::string &chipVersion);
    HcclResult GetOpInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo);
    HcclResult GetTilingRunInfo(nlohmann::json &opDescInfo,
        nlohmann::json &opTilingInfo, u64 inputCount, u64 outputCount, u32 valueDim,
            HcclDataType dataType, OpRunInfo &runInfo);
#ifndef HCCD
    HcclResult GetTilingRunInfo(nlohmann::json &opDescInfo, u64 inputCount, u64 outputCount, u32 valueDim,
            HcclDataType dataType, OpRunInfo &runInfo);
#endif

    HcclResult ExecuteKernelLaunch(const void *embVal, const void *refMat, const void *dst,
        void *stream, const void *axis, const void *stubFunc, OpRunInfo &runInfo, void *tilingDataDevPtr,
            const void *overFlowPtr);
    HcclResult GetopInfoIndex(std::string &opInfoIndex);

    HcclResult GetCoreNum();

    std::map<std::string, nlohmann::json> opInfoMap_;
    std::mutex tilingDataMutex_;
    std::mutex stubFuncMutex_;
    bool isInit_{};
    void *overFlowPtr_ = nullptr;
    int *axis_ = 0;
    HcclDataType indicesType_{ HCCL_DATA_TYPE_RESERVED };
};
} // namespace TbeReduce
#endif